#pragma once

// std
#include <string>
#include <unordered_map>
#include <map>
#include <functional>
#include <mutex>
#include <memory>
#include <chrono>
#include <thread>
#include <atomic>
#include <condition_variable>

// p2pms
#include "message.hpp"
#include "websocket_client.hpp"

// third-party
#include "minicoro/minicoro.h"

// webrtc
#include "rtc_base/thread.h"

namespace p2pms 
{

enum class CBType 
{
	Invalid, // 无效回调
	Normal,  // normal response
	Timeout, // timeout response
	Error    // call error
};

enum class DispatcherEventType
{
	CONNECTED,
	DISCONNECTED
};

enum class PendingType 
{
	Invalid,  // 无效状态
	Sync,     // 同步等待
	Async,    // 异步回调
	Coroutine // 协程等待
};

struct MsgResult 
{
    Message msg;
    CBType cbt;
};

using MessageHandler = std::function<void(const Message&, CBType)>;
using EventHandler = std::function<void(DispatcherEventType)>;

//==================================================================================================
// 
//==================================================================================================
class IMsgDispatcher 
{
public:
	virtual ~IMsgDispatcher() = default;

	// 初始化
	virtual bool Init(int timeout_ms, const EventHandler& event_handler) = 0;
	// 启动
	virtual bool Start(const std::string& uri) = 0;
	// 停止
	virtual void Stop() = 0;
	// 注册消息处理器
	virtual void Register(const std::string& cmd, const MessageHandler& handler) = 0;
	// 注销消息处理器
	virtual void Unregister(const std::string& cmd) = 0;
	// 发送消息，不等待响应，如果指定响应的命令，则会处理响应超时
	virtual bool SendMsg(const Message& msg, const std::string& res_cmd = "") = 0;
	// 同步发送消息，等待响应
	virtual MsgResult SendMsgSync(const Message& msg, const std::string& res_cmd) = 0;
	// 协程发送消息，等待响应
	virtual MsgResult SendMsgCoro(const Message& msg, const std::string& res_cmd) = 0;
};

//==================================================================================================
// 
//==================================================================================================
class WsMsgDispatcher : public IMsgDispatcher, public IWebSocketHandler
{
public:
	WsMsgDispatcher();
	~WsMsgDispatcher();

	// IMsgDispatcher
	bool Init(int timeout_ms, const EventHandler& event_handler) override;
	bool Start(const std::string& uri) override;
	void Stop() override;
	void Register(const std::string& cmd, const MessageHandler& handler) override;
	void Unregister(const std::string& cmd) override;
	bool SendMsg(const Message& msg, const std::string& res_cmd = "") override;
	MsgResult SendMsgSync(const Message& msg, const std::string& res_cmd) override;
	MsgResult SendMsgCoro(const Message& msg, const std::string& res_cmd) override;

private:
	// IWebSocketHandler接口实现
	virtual void OnConnected() override;
	virtual void OnDisconnected() override;
	virtual void OnMessage(const std::string& msg) override;
	
private:
	struct Pending 
	{
		PendingType type = PendingType::Invalid;
		Message request;
		std::string res_cmd;
		std::chrono::steady_clock::time_point expire;

		// For sync
		std::mutex mtx;
		std::condition_variable cv;

		// For coroutine
		mco_coro* co = nullptr;
		rtc::Thread* thread = nullptr;

		Message response;
		CBType cbt = CBType::Invalid;
	};

	struct PendingKey
	{
		std::string cmd;
		std::string rid;

		bool operator<(const PendingKey& other) const 
		{
			return cmd < other.cmd || (cmd == other.cmd && rid < other.rid);
		}

		std::string ToString() const 
		{
			return cmd + ":" + rid;
		}
	};

	void OnWsMessage(const std::string& msg);
	void TimeoutThreadFunc();

	std::mutex m_mutex;
	// key: cmd
	std::unordered_map<std::string, MessageHandler> m_cmd_handlers;
	std::map<PendingKey, std::shared_ptr<Pending>> m_pendings;

	int m_timeout_ms = 5000;
	EventHandler m_event_handler;
	std::atomic<bool> m_running{false};
	std::thread m_timeout_thread;

	// 你需要实现的WebSocketClient
	std::unique_ptr<WebSocketClient> m_ws_client;
};

} // namespace p2pms